/**
 * Copyright (C) 2011 Whisper Systems
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package com.appspot.eusms.crypto;

import android.util.Log;

import org.whispersystems.textsecure.crypto.MasterSecret;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.util.Arrays;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.Mac;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.ShortBufferException;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

/**
 * Class for streaming an encrypted MMS "part" off the disk.
 *
 * @author Moxie Marlinspike
 */

public class DecryptingPartInputStream extends FileInputStream {

    private static final int IV_LENGTH = 16;
    private static final int MAC_LENGTH = 20;

    private Cipher cipher;
    private Mac mac;

    private boolean done;
    private long totalDataSize;
    private long totalRead;
    private byte[] overflowBuffer;

    public DecryptingPartInputStream(File file, MasterSecret masterSecret) throws FileNotFoundException {
        super(file);
        try {
            if (file.length() <= IV_LENGTH + MAC_LENGTH)
                throw new FileNotFoundException("Part shorter than crypto overhead!");

            done = false;
            mac = initializeMac(masterSecret.getMacKey());
            cipher = initializeCipher(masterSecret.getEncryptionKey());
            totalDataSize = file.length() - cipher.getBlockSize() - mac.getMacLength();
            totalRead = 0;
        } catch (InvalidKeyException ike) {
            Log.w("EncryptingPartInputStream", ike);
            throw new FileNotFoundException("Invalid key!");
        } catch (NoSuchAlgorithmException e) {
            throw new AssertionError(e);
        } catch (InvalidAlgorithmParameterException e) {
            throw new AssertionError(e);
        } catch (NoSuchPaddingException e) {
            throw new AssertionError(e);
        } catch (IOException e) {
            Log.w("EncryptingPartInputStream", e);
            throw new FileNotFoundException("IOException while reading IV!");
        }
    }

    @Override
    public int read(byte[] buffer) throws IOException {
        return read(buffer, 0, buffer.length);
    }

    @Override
    public int read(byte[] buffer, int offset, int length) throws IOException {
        if (totalRead != totalDataSize)
            return readIncremental(buffer, offset, length);
        else if (!done)
            return readFinal(buffer, offset, length);
        else
            return -1;
    }

    private int readFinal(byte[] buffer, int offset, int length) throws IOException {
        try {
            int flourish = cipher.doFinal(buffer, offset);
            //			mac.update(buffer, offset, flourish);

            byte[] ourMac = mac.doFinal();
            byte[] theirMac = new byte[mac.getMacLength()];
            readFully(theirMac);

            if (!Arrays.equals(ourMac, theirMac))
                throw new IOException("MAC doesn't match! Potential tampering?");

            done = true;
            return flourish;
        } catch (IllegalBlockSizeException e) {
            Log.w("EncryptingPartInputStream", e);
            throw new IOException("Illegal block size exception!");
        } catch (ShortBufferException e) {
            Log.w("EncryptingPartInputStream", e);
            throw new IOException("Short buffer exception!");
        } catch (BadPaddingException e) {
            Log.w("EncryptingPartInputStream", e);
            throw new IOException("Bad padding exception!");
        }
    }

    private int readIncremental(byte[] buffer, int offset, int length) throws IOException {
        int readLength = 0;
        if (null != overflowBuffer) {
            if (overflowBuffer.length > length) {
                System.arraycopy(overflowBuffer, 0, buffer, offset, length);
                overflowBuffer = Arrays.copyOfRange(overflowBuffer, length, overflowBuffer.length);
                return length;
            } else if (overflowBuffer.length == length) {
                System.arraycopy(overflowBuffer, 0, buffer, offset, length);
                overflowBuffer = null;
                return length;
            } else {
                System.arraycopy(overflowBuffer, 0, buffer, offset, overflowBuffer.length);
                readLength += overflowBuffer.length;
                offset += readLength;
                length -= readLength;
                overflowBuffer = null;
            }
        }

        if (length + totalRead > totalDataSize)
            length = (int) (totalDataSize - totalRead);

        byte[] internalBuffer = new byte[length];
        int read = super.read(internalBuffer, 0, internalBuffer.length <= cipher.getBlockSize() ? internalBuffer.length : internalBuffer.length - cipher.getBlockSize());
        totalRead += read;

        try {
            mac.update(internalBuffer, 0, read);

            int outputLen = cipher.getOutputSize(read);

            if (outputLen <= length) {
                readLength += cipher.update(internalBuffer, 0, read, buffer, offset);
                return readLength;
            }

            byte[] transientBuffer = new byte[outputLen];
            outputLen = cipher.update(internalBuffer, 0, read, transientBuffer, 0);
            if (outputLen <= length) {
                System.arraycopy(transientBuffer, 0, buffer, offset, outputLen);
                readLength += outputLen;
            } else {
                System.arraycopy(transientBuffer, 0, buffer, offset, length);
                overflowBuffer = Arrays.copyOfRange(transientBuffer, length, outputLen);
                readLength += length;
            }
            return readLength;
        } catch (ShortBufferException e) {
            throw new AssertionError(e);
        }
    }

    private Mac initializeMac(SecretKeySpec key) throws NoSuchAlgorithmException, InvalidKeyException {
        Mac hmac = Mac.getInstance("HmacSHA1");
        hmac.init(key);

        return hmac;
    }

    private Cipher initializeCipher(SecretKeySpec key)
            throws InvalidKeyException, InvalidAlgorithmParameterException,
            NoSuchAlgorithmException, NoSuchPaddingException, IOException {
        Cipher cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
        IvParameterSpec iv = readIv(cipher.getBlockSize());
        cipher.init(Cipher.DECRYPT_MODE, key, iv);

        return cipher;
    }

    private IvParameterSpec readIv(int size) throws IOException {
        byte[] iv = new byte[size];
        readFully(iv);

        mac.update(iv);
        return new IvParameterSpec(iv);
    }

    private void readFully(byte[] buffer) throws IOException {
        int offset = 0;

        for (; ; ) {
            int read = super.read(buffer, offset, buffer.length - offset);

            if (read + offset < buffer.length) offset += read;
            else return;
        }
    }
}
